from __future__ import annotations
from typing import List
from ...dto.image import ImageGenerationRequest
from threading import Thread, Event
from queue import Queue
from PIL import Image
import asyncio


class SingletonSimulator:  # 资源单例(持有模型, 队列等资源)
    __slots__ = ["pipe", "queue", "max_batch_per_round", "wait_time", "_thread"]
    _instance: SingletonSimulator = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(SingletonSimulator, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if hasattr(self, "queue"):
            return
        self.max_batch_per_round = 10
        self.wait_time = 0.2 # 200ms
        self.queue: Queue[asyncio.Future, ImageGenerationRequest] = Queue()
        self._thread: Thread = Thread(target=self.run, daemon=True)
        self._thread.start()

    def run(self):
        pulse_event = Event() # 攒批脉冲事件
        while True:
            pulse_event.wait(timeout=self.wait_time)
            match self.queue.qsize():
                case 0: # 仅该情况下会被event阻塞
                    continue 
                case 1:
                    future, request = self.queue.get()
                    generate_single(future, request)
                case x:
                    futures, requests = zip(*[self.queue.get() for _ in range(min(x, self.max_batch_per_round))])
                    generate_batch(futures, requests)
            pulse_event.clear()
            self.queue.task_done()


def generate_single(future: asyncio.Future, request: ImageGenerationRequest):
    width, height = map(int, request.size.split("x"))
    img = Image.new("RGB", (width, height), color="black")
    future.set_result(img)


def generate_batch(
    futures: List[asyncio.Future],
    requests: List[ImageGenerationRequest],
):
    """
    批量处理图像生成请求
    """
    prompts = []
    images = []
    for request in requests:
        prompts.append(request.prompt)
        width, height = map(int, request.size.split("x"))
    # 模拟生成图像(这里可以替换为实际的批量图像生成逻辑)
    images = [Image.new("RGB", (width, height), color="black") for _ in prompts]
    assert len(images) == len(futures), "生成的图像数量与请求数量不匹配"
    for future, image in zip(futures, images):
        future.set_result(image)


def concurrent_simulate(future: asyncio.Future, request: ImageGenerationRequest):
    """
    模拟图像生成请求的并发处理, 接收到请求后生成单例持有资源
    """
    singleton = SingletonSimulator()
    singleton.queue.put((future, request))
